#include "cmdlib.h"
#include "mathlib.h"
#include "bspfile.h"
#include "polylib.h"
#include "imagelib.h"
#include "threads.h"
#include "scriplib.h"
#include "progress.h"

#include "shaders.h"
#include "mesh.h"

#ifdef _WIN32
#include "../libs/pakstuff.h"
#endif

void CountLightmaps( void );
void RechartLightmaps(void (*getDimensions)(int surfNum, int *w, int *h));


static float lightPackTolerance = 1.0f;
static float lightPackScale = 32.0f;

extern int numLightBytes;
extern int numLightTexels;
static int numFILightmaps;

#define COEFS_PER_ELEMENT (MAP_BASIS_VECTORS * 3)

typedef struct
{
	int w, h;
	float *coefs;
} surfLightmap_t;

surfLightmap_t *surfLightmaps;

int FileCoefLocation(int x, int y, int band)
{
	int headerSize = sizeof(int)*2 + (((LIGHTMAP_COMPOSITE_WIDTH*LIGHTMAP_COMPOSITE_HEIGHT) + 7) / 8);

	x += (band % 3) * LIGHTMAP_WIDTH;
	y += (band / 3) * LIGHTMAP_HEIGHT;

	return headerSize + (x + y*LIGHTMAP_COMPOSITE_WIDTH) * sizeof(float) * 3;
}

float *LMCoefsAtCoord(surfLightmap_t *slm, int x, int y, int band)
{
	if(((y*slm->w + x) * MAP_BASIS_VECTORS + band)  < 0 ||
		((y*slm->w + x) * MAP_BASIS_VECTORS + band) >= slm->w * slm->h * COEFS_PER_ELEMENT)
		printf("Smashed coef\n");
	return slm->coefs + ((y*slm->w + x) * MAP_BASIS_VECTORS + band) * 3;
}

void LoadLightmapsFromFI( const char *baseName )
{
	int i, j, x, y, band;
	int biggest = 0;
	float *baseCoefs;
	FILE *f;
	char fiName[1024];

	// Init surface lightmaps list
	surfLightmaps = malloc(numDrawSurfaces * sizeof(surfLightmap_t));
	memset(surfLightmaps, 0, numDrawSurfaces * sizeof(surfLightmap_t));

	for(i=0;i<numFILightmaps;i++)
	{
		sprintf(fiName, "%s_lm%i.fi", baseName, i);

		f = fopen(fiName, "rb");

		for(j=0;j<numDrawSurfaces;j++)
		{
			dsurface_t *dsurf = drawSurfaces + j;

			if(dsurf->lightmapNum == i*MAP_BASIS_VECTORS)
			{
				surfLightmap_t *slm = surfLightmaps + j;

				slm->w = dsurf->lightmapWidth;
				slm->h = dsurf->lightmapHeight;

				if(slm->w * slm->h > biggest)
				{
					printf("%i: %i x %i\n", j, slm->w, slm->h);
					biggest = slm->w * slm->h;
				}

				// Found the file for this lightmap
				slm->coefs = baseCoefs = malloc(slm->w * slm->h * sizeof(float) * COEFS_PER_ELEMENT);

				for(band=0;band<MAP_BASIS_VECTORS;band++)
				{
					for(y=0;y<slm->h;y++)
					{
						// Seek to the coefficients in this band
						fseek(f, FileCoefLocation(dsurf->lightmapX, dsurf->lightmapY + y, band), SEEK_SET);
						for(x=0;x<slm->w;x++)
						{
							// Read this coef
							fread(LMCoefsAtCoord(slm, x, y, band), sizeof(float) * 3, 1, f);
						}
					}
				}
			}
		}

		fclose(f);
	}
}


void ResaveLightmapsToFI( const char *baseName )
{
	int i, j, x, y, band;
	float *baseCoefs;
	FILE *f;
	char fiName[1024];
	float zeros[3];
	int numZeros;

	memset(zeros, 0, sizeof(zeros));

	// Init surface lightmaps list
	for(i=0;i<numFILightmaps;i++)
	{
		sprintf(fiName, "%s_lm%i.fi", baseName, i);

		f = fopen(fiName, "rb+");

		// Clear out existing image data
		fseek(f, FileCoefLocation(0, 0, 0), SEEK_SET);

		zeros[0] = 255.0f;
		zeros[1] = 0.0f;
		zeros[2] = 0.0f;

		numZeros = LIGHTMAP_WIDTH * LIGHTMAP_HEIGHT * COEFS_PER_ELEMENT * sizeof(float);
		while(numZeros > sizeof(zeros))
		{
			fwrite(zeros, sizeof(zeros), 1, f);
			numZeros -= sizeof(zeros);
		}
		fwrite(zeros, numZeros, 1, f);

		for(j=0;j<numDrawSurfaces;j++)
		{
			dsurface_t *dsurf = drawSurfaces + j;

			if(dsurf->lightmapNum == i*MAP_BASIS_VECTORS)
			{
				surfLightmap_t *slm = surfLightmaps + j;

				// This lightmap's in this file
				for(band=0;band<MAP_BASIS_VECTORS;band++)
				{
					for(y=0;y<slm->h;y++)
					{
						// Seek to the coefficients in this band
						fseek(f, FileCoefLocation(dsurf->lightmapX, dsurf->lightmapY + y, band), SEEK_SET);
						for(x=0;x<slm->w;x++)
						{
							// Read this coef
							fwrite(LMCoefsAtCoord(slm, x, y, band), sizeof(float) * 3, 1, f);
						}
					}
				}
			}
		}

		fclose(f);
	}
}



void FreeSurfLightmap(surfLightmap_t *slm)
{
	free(slm->coefs);
	slm->coefs = NULL;
	free(slm);
}

surfLightmap_t *ResampleLightmap(surfLightmap_t *base, int newW, int newH)
{
	float *edges[2];
	float *dest;
	int x, y;
	float *coefs;
	float *scan;
	float blX, blY;
	int left, right, top, bottom;
	float ix, iy;
	int i, j;
	int inVarDim, outVarDim, fixedDim;
	int fixedC, varC;
	int startC, endC;

	float minDivisor;
	float actualDivisor;
	float startFactor, endFactor;
	float startPoint, endPoint, midPoint, interpolator;

	int inVarStride, inFixedStride, outVarStride, outFixedStride;

	surfLightmap_t *result;

	// Only resample one dimension at a time
	if(newW != base->w && newH != base->h)
	{
		surfLightmap_t *temp1, *temp2;
		temp1 = ResampleLightmap(base, newW, base->h);
		temp2 = ResampleLightmap(temp1, newW, newH);
		FreeSurfLightmap(temp1);
		return temp2;
	}

	// index = fixedAxis * fixedStride + varAxis * varStride
	if(newW != base->w)
	{
		// Horizontal rescale
		inVarStride = 1;
		outVarStride = 1;
		inFixedStride = base->w;
		outFixedStride = newW;
		inVarDim = base->w;
		outVarDim = newW;
		fixedDim = base->h;
	}
	else
	{
		// Vertical rescale
		inVarStride = base->w;
		outVarStride = newW;
		inFixedStride = 1;
		outFixedStride = 1;

		inVarDim = base->h;
		outVarDim = newH;
		fixedDim = base->w;
	}

	result = malloc(sizeof(surfLightmap_t));

	result->coefs = coefs = malloc(sizeof(float) * COEFS_PER_ELEMENT * newW * newH);
	result->w = newW;
	result->h = newH;

	for(fixedC=0;fixedC<fixedDim;fixedC++)
	{
		for(varC=0;varC<outVarDim;varC++)
		{
			if(outVarDim >= inVarDim)
			{
				// Magnify (or no change)
				if(outVarDim == 1)
					midPoint = 0.0f;
				else
					midPoint = ((float)varC) * (float)(inVarDim-1) / (float)(outVarDim-1);
				startPoint = floor(midPoint);
				endPoint = ceil(midPoint);
				interpolator = midPoint - startPoint;

				edges[0] = base->coefs + (fixedC*inFixedStride + ((int)startPoint)*inVarStride) * COEFS_PER_ELEMENT;
				edges[1] = base->coefs + (fixedC*inFixedStride + ((int)endPoint)*inVarStride) * COEFS_PER_ELEMENT;

				dest = coefs + (fixedC*outFixedStride + varC*outVarStride) * COEFS_PER_ELEMENT;
				for(i=0;i<COEFS_PER_ELEMENT;i++)
					dest[i] = edges[0][i] * (1.0f - interpolator) + edges[1][i] * interpolator;
			}
			else
			{
				float tempCoefs[COEFS_PER_ELEMENT];

				// To keep edge values the same, edges are always directly copied from their original values
				// only middle values are filtered
				if(varC == 0)
				{
					scan = base->coefs + fixedC*inFixedStride * COEFS_PER_ELEMENT;
					for(i=0;i<COEFS_PER_ELEMENT;i++)
						tempCoefs[i] = scan[i];
					actualDivisor = 1.0f;
				}
				else if(varC == outVarDim-1)
				{
					scan = base->coefs + ((inVarDim-1)*inVarStride + fixedC*inFixedStride) * COEFS_PER_ELEMENT;
					for(i=0;i<COEFS_PER_ELEMENT;i++)
						tempCoefs[i] = scan[i];
					actualDivisor = 1.0f;
				}
				else
				{
					startPoint = ((float)(varC-1) * (float)(inVarDim-2) / (float)(outVarDim-2)) + 1;
					endPoint = ((float)(varC) * (float)(inVarDim-2) / (float)(outVarDim-2)) + 1;

					// Minify
					//startPoint = ((float)varC) * ((float)inVarDim) / ((float)outVarDim);
					//endPoint = ((float)varC+1.0f) * ((float)inVarDim) / ((float)outVarDim);

					// All pixels from startC to endC are included
					startC = floor(startPoint);
					endC = ceil(endPoint - 1.0f);

					startFactor = floor(startPoint + 1.0f) - startPoint;
					endFactor = 1.0f - (ceil(endPoint) - endPoint);

					for(i=0;i<COEFS_PER_ELEMENT;i++)
						tempCoefs[i] = 0.0f;

					// Add all of them in
					actualDivisor = 0.0f;
					for(i=startC;i<=endC;i++)
					{
						if(i == startC) interpolator = startFactor;
						else if(i == endC) interpolator = endFactor;
						else interpolator = 1.0f;

						actualDivisor += interpolator;

						scan = base->coefs + (i*inVarStride + fixedC*inFixedStride) * COEFS_PER_ELEMENT;
						for(j=0;j<COEFS_PER_ELEMENT;j++)
							tempCoefs[j] += scan[j] * interpolator;
					}
				}

				// Write out
				dest = coefs + (fixedC*outFixedStride + varC*outVarStride) * COEFS_PER_ELEMENT;

				for(i=0;i<COEFS_PER_ELEMENT;i++)
					dest[i] = tempCoefs[i] / actualDivisor;
			}
		}
	}

	return result;
}


// Returns the difference between two lightmaps
float LightmapDiffMSE(surfLightmap_t *a, surfLightmap_t *b)
{
	int numCoefs = a->w*a->h * COEFS_PER_ELEMENT;
	float totalError = 0.0f;
	float diff;
	int i, band;
	float bandAverages[MAP_BASIS_VECTORS];

	// Find the average signal in the base chart
	for(i=0;i<MAP_BASIS_VECTORS;i++)
		bandAverages[i] = 0.0f;

	for(i=0;i<numCoefs;i++)
	{
		band = (i/3) % MAP_BASIS_VECTORS;
		bandAverages[band] += a->coefs[i];
	}

	for(i=0;i<MAP_BASIS_VECTORS;i++)
	{
		if(!bandAverages[i])
			bandAverages[i] = 1.0f;
		else
			bandAverages[i] /= (float)(numCoefs / MAP_BASIS_VECTORS);
	}

	for(i=0;i<numCoefs;i++)
	{
		band = (i/3) % MAP_BASIS_VECTORS;
		diff = (a->coefs[i] - b->coefs[i]) / bandAverages[band];
		totalError += diff*diff;
	}

	return totalError / (float)(a->w*a->h);
}

float LargestGradientForCoefs(const float *coefs[COEFS_PER_ELEMENT], float prevLargest)
{
	int c;

	// Determine largest gradient
	/*
	for(c=0;c<COEFS_PER_ELEMENT;c++)
	{
		float avg, coefLength, grad;

		coefLength = coefs[2][c]-coefs[1][c];
		coefLength = sqrt(coefLength*coefLength + baseWeight*baseWeight);

		avg = (coefs[0][c] + coefs[2][c]) * 0.5f;

		grad = (avg - coefs[1][c]) / coefLength;

		if(grad > largestGradient[0])
			largestGradient[0] = grad;
	}
	*/

	for(c=0;c<COEFS_PER_ELEMENT;c++)
	{
		float grad;
		float coef1 = coefs[0][c] / lightPackTolerance;
		float coef2 = coefs[2][c] / lightPackTolerance;

		if(coef1 <= 1.0f) coef1 = 1.0000f;
		if(coef2 <= 1.0f) coef2 = 1.0000f;

		grad = fabs(log(coefs[2][c]*coefs[2][c]) - log(coefs[1][c]*coefs[1][c])) / lightPackScale;
		if(grad > 1.0f) grad = 1.0f;

		if(grad > prevLargest)
			prevLargest = grad;
	}
	return prevLargest;
}


void CompactLightmapThread(int surfNum)
{
	int newW, newH;
	surfLightmap_t *best = NULL;
	float bestErr;
	surfLightmap_t *base;
	surfLightmap_t *challenger = NULL;
	surfLightmap_t *challengerResized = NULL;
	dsurface_t *surf = drawSurfaces + surfNum;
	float err;
	qboolean won;
	int x, y, diag;
	float largestGradient[2];
	float baseWeight = 255.0f;

	int diagonalDimension;

	if(surf->lightmapNum < 0)
		return;		// No lightmap

	base = surfLightmaps + surfNum;

	// Use a diagonal pattern so small hits within bounds are likely to get hit first
	diagonalDimension = surf->lightmapWidth + surf->lightmapHeight;

	/*
	for(x=0;x<=diagonalDimension;x++)
	{
		for(diag=0;diag<=diagonalDimension-x;diag++)
		{
			newW = x - diag;
			newH = diag;

			if(newW < 2 || newW > surf->lightmapWidth ||
				newH < 2 || newH > surf->lightmapHeight)
				continue;

			if(newH == surf->lightmapHeight && newW == surf->lightmapWidth)
				continue;	// No resampling at all

			if(best && newW * newH > best->w*best->h)
				continue;	// Not a size reduction

			// Get the challenger size
			challenger = ResampleLightmap(base, newW, newH);
			challengerResized = ResampleLightmap(challenger, surf->lightmapWidth, surf->lightmapHeight);

			err = LightmapDiffMSE(base, challengerResized);
			FreeSurfLightmap(challengerResized);

			won = qfalse;
			if(err < lightPackTolerance)
			{
				if(!best)
					won = qtrue;	// Always win if there's no winner
				else if(newW * newH < best->w*best->h)
					won = qtrue;	// Smaller
				else if(err < bestErr)
					won = qtrue;	// Higher-quality
			}

			if(won)
			{
				// Become best
				if(best) FreeSurfLightmap(best);
				best = challenger;
				bestErr = err;
			}
			else
			{
				// Free challenger
				FreeSurfLightmap(challenger);
			}
		}
	}
	*/

	largestGradient[0] = largestGradient[1] = 0.0f;
	if(base->w >= 3)
	{
		for(y=0;y<base->h;y++)
		{
			for(x=1;x<base->w-1;x++)
			{
				float *coefs[3];
				int c;

				coefs[0] = base->coefs + (x-1 + y*base->w) * COEFS_PER_ELEMENT;
				coefs[1] = base->coefs + (x + y*base->w) * COEFS_PER_ELEMENT;
				coefs[2] = base->coefs + (x+1 + y*base->w) * COEFS_PER_ELEMENT;

				// Determine largest gradient
				/*
				for(c=0;c<COEFS_PER_ELEMENT;c++)
				{
					float avg, coefLength, grad;

					coefLength = coefs[2][c]-coefs[1][c];
					coefLength = sqrt(coefLength*coefLength + baseWeight*baseWeight);

					avg = (coefs[0][c] + coefs[2][c]) * 0.5f;

					grad = (avg - coefs[1][c]) / coefLength;

					if(grad > largestGradient[0])
						largestGradient[0] = grad;
				}
				*/

				largestGradient[0] = LargestGradientForCoefs(coefs, largestGradient[0]);
			}
		}
	}

	if(base->h >= 3)
	{
		for(y=1;y<base->h-1;y++)
		{
			for(x=0;x<base->w;x++)
			{
				float *coefs[3];
				int c;

				coefs[0] = base->coefs + (x + (y-1)*base->w) * COEFS_PER_ELEMENT;
				coefs[1] = base->coefs + (x + y*base->w) * COEFS_PER_ELEMENT;
				coefs[2] = base->coefs + (x + (y+1)*base->w) * COEFS_PER_ELEMENT;

				// Determine largest gradient
				/*
				for(c=0;c<COEFS_PER_ELEMENT;c++)
				{
					float avg, coefLength, grad;

					coefLength = coefs[2][c]-coefs[1][c];
					coefLength = sqrt(coefLength*coefLength + baseWeight*baseWeight);

					avg = (coefs[0][c] + coefs[2][c]) * 0.5f;

					grad = (avg - coefs[1][c]) / coefLength;

					if(grad > largestGradient[1])
						largestGradient[1] = grad;
				}
				*/
				largestGradient[1] = LargestGradientForCoefs(coefs, largestGradient[1]);
			}
		}
	}

	_printf("%f / %f\n", largestGradient[0], largestGradient[1]);
	if(lightPackTolerance < largestGradient[0]) newW = base->w;
	else newW = base->w * largestGradient[0];

	if(lightPackTolerance < largestGradient[1]) newH = base->h;
	else newH = base->h * largestGradient[1];

	// Round up
	if(newW <= 0) newW = LIGHTMAP_CLUSTER_SIZE;
	if(newH <= 0) newH = LIGHTMAP_CLUSTER_SIZE;

	newW += LIGHTMAP_CLUSTER_SIZE - 1;
	newW -= newW % LIGHTMAP_CLUSTER_SIZE;
	newH += LIGHTMAP_CLUSTER_SIZE - 1;
	newH -= newH % LIGHTMAP_CLUSTER_SIZE;

	if(newW != surf->lightmapWidth || newH != surf->lightmapHeight)
	{
		surfLightmap_t *lm = ResampleLightmap(base, newW, newH);
		free(base->coefs);
		memcpy(base, lm, sizeof(surfLightmap_t));
	}

	_printf("Highest gradient at %i(%i, %i): %f %f  (Rescale from %i x %i to %i x %i)\n\n", surf->lightmapNum, surf->lightmapX, surf->lightmapY, largestGradient[0], largestGradient[1], surf->lightmapWidth, surf->lightmapHeight, newW, newH);
}

void DumpTest(surfLightmap_t *slm)
{
	int nCoefs, i;
	unsigned char *bytes;
	FILE *f;

	nCoefs = slm->w*slm->h*COEFS_PER_ELEMENT;

	bytes = malloc(nCoefs);
	for(i=0;i<nCoefs;i++)
	{
		float coef = slm->coefs[i] * 0.5 + 127.5f;

		if(coef < 0) bytes[i] = 0;
		else if(coef > 255) bytes[i] = 255;
		else bytes[i] = coef;
	}

	f = fopen("test.raw", "wb");
	fwrite(bytes, 1, nCoefs, f);
	fclose(f);

	free(bytes);
}

static void GetSurfDimensions(int surfNum, int *w, int *h)
{
	*w = surfLightmaps[surfNum].w;
	*h = surfLightmaps[surfNum].h;
}

/*
========
PackLightsMain

========
*/
int PackLightsMain (int argc, char **argv) {
	int			i;
	double		start, end;
	const char	*value;
	char		source[1024];

	_printf ("----- Lightmap Compression ----\n");

	Progress_Init();

	verbose = qfalse;

	for (i=1 ; i<argc ; i++) {
		if (!strcmp(argv[i],"-tempname"))
		{
			i++;
		} else if (!strcmp(argv[i],"-tol")) {
			lightPackTolerance = atof(argv[i+1]);
			_printf ("pack error tolerance set to %f\n", lightPackTolerance);
			i++;
		} else {
			break;
		}
	}

	if (i != argc - 1) {
		_printf("usage: q3map -packlights [-<switch> [-<switch> ...]] <mapname>\n"
				"\n"
				"Switches:\n"
				"   tol            = tolerance level\n");
		exit(0);
	}

	start = I_FloatTime ();

	SetQdirFromPath (argv[i]);	

#ifdef _WIN32
	InitPakFile(gamedir, NULL);
#endif

	strcpy (source, ExpandArg(argv[i]));
	StripExtension (source);
	DefaultExtension (source, ".bsp");

	LoadShaderInfo();

	_printf ("reading %s\n", source);

	LoadBSPFile (source);

	CountLightmaps();
	numFILightmaps = numLightBytes / LIGHTMAP_WIDTH / LIGHTMAP_HEIGHT / 3 / MAP_BASIS_VECTORS;
	_printf ("%i source lightmaps, %i lightmap texels\n", numFILightmaps, numLightTexels);

	StripExtension (source);

	LoadLightmapsFromFI(source);

	Progress_SetMessage("Recharting lightmaps...");
	RunThreadsOnIndividual(numDrawSurfaces, qtrue, CompactLightmapThread);

	RechartLightmaps(GetSurfDimensions);

	CountLightmaps();
	numFILightmaps = numLightBytes / LIGHTMAP_WIDTH / LIGHTMAP_HEIGHT / 3 / MAP_BASIS_VECTORS;
	_printf ("%i result lightmaps, %i lightmap texels\n", numFILightmaps, numLightTexels);

	ResaveLightmapsToFI(source);

	StripExtension (source);
	DefaultExtension (source, ".bsp");
	_printf ("writing %s\n", source);
	WriteBSPFile (source);

	end = I_FloatTime ();
	_printf ("%5.0f seconds elapsed\n", end-start);

	/*
	for(i=0;i<numLightBytes;i+= LIGHTMAP_BYTESIZE)
	{
		char filename[1024];
		sprintf(filename, "lightmaptest%4i.tga", i /LIGHTMAP_BYTESIZE);
		WriteTGA(filename, lightBytes + i, LIGHTMAP_WIDTH, LIGHTMAP_HEIGHT, 3);
	}
	*/

	return 0;
}

